package com.me.study.javaCore.ssl;

import com.me.common.util.StringUtils;
import org.apache.commons.lang3.RandomStringUtils;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLServerSocket;
import javax.net.ssl.TrustManagerFactory;
import java.io.*;
import java.net.Socket;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.KeyStore;

/**
 * TODO 实际上，服务端的 KeyManager、客户端的 TrustManager 是成对的，客户端的 KeyManager、服务端的 TrustManager 同理。此处偷懒使用同一份证书才有巧合
 *
 * Java SSL 认证：服务端
 *
 * 1，生成证书：.\keytool -genkeypair -alias serverkey -keypass 123456 -keyalg RSA -keystore kserver.jks -storepass 123456 -dname "CN=, OU=, O=, L=, ST=, C="
 * 2，提取公钥：.\keytool -exportcert -alias serverkey -keystore kserver.jks -file server.crt
 * 3，生成（提供给客户端的）含有服务端公钥证书的密钥存储库：.\keytool -importcert -alias serverkey -file server.crt -keystore tclient.jks
 *
 * 服务端证书：kserver.jks
 * 服务端公钥：server.crt
 * 提供给客户端的信任证书（含服务的公钥）：tclient.jks
 *
 * @see <a href="https://www.cnblogs.com/franson-2016/p/5557259.html">SSLSocket 实现服务端和客户端双向认证的例子</a>
 * @see <a href="https://blog.csdn.net/fw0124/article/details/41013333">Java SSL/TLS 编程代码实例-双向认证</a>
 *
 * @author ME
 * @date 2021/8/9
 */
public class SSLServer {

    /**
     * TODO 预定义的消息结束符，缺了的话，字节流会一直阻塞无法终止
     */
    public static final String EOF = "eof";
    public static final String END = EOF + System.lineSeparator();

    private static final int DEFAULT_PORT = 7777;
    private static final String SERVER_KEY_STORE_PASSWORD = "123456";
    private static final String SERVER_KEY_PASSWORD = "123456";
    private static final String SERVER_TRUST_KEY_STORE_PASSWORD = "123456";
    private SSLServerSocket serverSocket;

    /**
     * 启动程序
     *
     * @param args
     */
    public static void main(String[] args) {
        SSLServer server = new SSLServer();
        server.init(true);
        server.start();
    }

    /**
     * 监听 SSL Server Socket，接受客户端的消息，并且返回客户端指定消息：
     * TODO 1，如果使用字节流读取报文，则需要指定协议（关键是报文长度），要不然读完一轮之后会阻塞在下次读取，无法终止。。。
     *  2，也可以使用 BufferedReader、DataInputStream 的 readLine()，然后拼上字鼎固 结束符
     */
    public void start() {
        if (serverSocket == null) {
            System.out.println("ERROR");
            return;
        }
        while (true) {
            try {
                new Thread(new ServiceHandler(serverSocket.accept())).start();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private static class ServiceHandler implements Runnable {

        private final Socket socket;

        public ServiceHandler(Socket socket) {
            this.socket = socket;
        }

        @Override
        public void run() {
            try {
                DataInputStream input = new DataInputStream(new BufferedInputStream(socket.getInputStream()));
                // 循环读取
                StringBuilder msg = new StringBuilder();
                String temp;
                int index;
                while (StringUtils.isNotEmpty(temp = input.readUTF())) {
                    // 遇到 eof 时就结束接收
                    if ((index = temp.indexOf(EOF)) != -1) {
                        msg.append(temp, 0, index);
                        break;
                    }
                    msg.append(temp);
                }
                System.out.println(msg);

                // 写入响应内容
                DataOutputStream out = new DataOutputStream(new BufferedOutputStream(socket.getOutputStream()));
                out.writeUTF("Server Echo: " + RandomStringUtils.random(3, true, false));
                out.writeUTF(END);
                out.flush();
            } catch (Exception e) {
                System.out.println(e.getMessage());
            } finally {
                try {
                    socket.close();
                } catch (IOException e) {
                    System.err.println("socket close failed！");
                }
            }
        }
    }

    /**
     * ssl连接的重点：
     * 初始化 SSLServerSocket
     * KeyManager：导入服务端私钥 KeyStore，负责提供证书和私钥
     * TrustManager：导入服务端需要信任的 TrustStore（客户端提供的公钥证书），验证对方证书的合法性，并提取公钥
     */
    public void init(boolean isTwoWayAuth) {
        try {
            KeyStore ks = KeyStore.getInstance("JKS");
            ks.load(Files.newInputStream(Paths.get("testJDK8/src/main/java/com/me/study/javaCore/ssl/kserver.jks")),
                    SERVER_KEY_STORE_PASSWORD.toCharArray());
            // KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
            KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            kmf.init(ks, SERVER_KEY_PASSWORD.toCharArray());

            KeyStore tks = KeyStore.getInstance("JKS");
            // password 可不填，填写则必须为正确的值
            tks.load(Files.newInputStream(Paths.get("testJDK8/src/main/java/com/me/study/javaCore/ssl/tclient.jks")),
                    SERVER_TRUST_KEY_STORE_PASSWORD.toCharArray());
            TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            tmf.init(tks);
            System.out.println(tmf.getAlgorithm());

            SSLContext ctx = SSLContext.getInstance("SSL");
            if (isTwoWayAuth) {// 双向
                ctx.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);
            }
            else {// 单向
                ctx.init(kmf.getKeyManagers(), null, null);
            }

            serverSocket = (SSLServerSocket) ctx.getServerSocketFactory().createServerSocket(DEFAULT_PORT);
            // 强制要求客户端认证。如果客户端未提供有效的证书，SSL/TLS 握手将失败，连接将被拒绝。
            serverSocket.setNeedClientAuth(true);
            // 希望客户端进行认证，但不强制要求。
            // serverSocket.setWantClientAuth(true);

            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                try {
                    serverSocket.close();
                } catch (IOException e) {
                    System.out.println("serverSocket close failed!");
                }
            }, "shutdown"));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}
